import numpy as np
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, HttpUrl
import base64
import cv2
from io import BytesIO
from PIL import Image
from ultralytics import YOLO
import uvicorn

modelPath = 'E:\\newDownload\\yolo12x.pt'
model = YOLO(model=modelPath)
app = FastAPI()


# 定义请求模型
class Detect(BaseModel):
    imageBase64: str


# 定义响应模型
class DetectResult(BaseModel):
    resultImage: str
    detectionCount: int
    detectedObjects: list


@app.get("/")
def read_root():
    return {"Hello": "World"}


@app.post('/detect', response_model=DetectResult)
async def detect(detectObj: Detect):
    try:
        # 转换base64为OpenCV图像
        image = base64_to_cv2(detectObj.imageBase64)
        # 使用YOLO进行目标检测
        results = model(
            image,
            device='cpu'
        )
        # 获取第一个结果（因为只处理单张图片）
        result = results[0]
        # 在图像上绘制检测结果
        annotated_image = result.plot()
        # 统计检测到的对象
        detection_count = len(result.boxes)
        detected_objects = []
        # 提取检测结果信息
        for box in result.boxes:
            obj = {
                "class_id": int(box.cls),
                "class_name": result.names[int(box.cls)],
                "confidence": float(box.conf),
                "bbox": box.xyxy[0].tolist()  # [x1, y1, x2, y2]
            }
            detected_objects.append(obj)
        # 转换处理后的图像为base64
        processed_base64 = cv2_to_base64(annotated_image)
        return {
            "resultImage": processed_base64,
            "detectionCount": detection_count,
            "detectedObjects": detected_objects
        }

    except Exception as e:
        raise HTTPException(status_code=400, detail=f"图片处理失败: {str(e)}")


def base64_to_cv2(image_base64: str) -> np.ndarray:
    """将base64字符串转换为OpenCV图像"""
    # 移除可能的数据URL前缀
    if ',' in image_base64:
        image_base64 = image_base64.split(',')[1]
    # 解码base64
    image_data = base64.b64decode(image_base64)
    image = Image.open(BytesIO(image_data))
    # 转换为OpenCV格式 (BGR)
    return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)


def cv2_to_base64(image: np.ndarray) -> str:
    """将OpenCV图像转换为base64字符串"""
    # 转换为RGB
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 转换为PIL图像
    pil_img = Image.fromarray(image_rgb)
    # 转换为base64
    buffered = BytesIO()
    pil_img.save(buffered, format="JPEG")
    return f"data:image/jpeg;base64,{base64.b64encode(buffered.getvalue()).decode('utf-8')}"


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)
